import requests
import os
import base64
from typing import Dict, List, Tuple
import uuid
import config
from PIL import Image
import io


class ImageGenerator:
    def __init__(self):
        self.api_key = config.IMAGE_GEN_API_KEY
        self.api_url = "https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json"
        }

    def generate(self, content: Dict, custom_style: str = None) -> List[str]:
        """Generate images based on poster content and custom style."""
        image_paths = []

        # 使用用户自定义风格作为主要提示词
        if custom_style:
            # 直接使用用户的自定义风格描述作为提示词
            visualization_prompts = [custom_style]
        else:
            # 原有的提取逻辑作为备选
            visualization_prompts = []
            # If there are specific visualization suggestions in the content
            if "visualizations" in content and isinstance(content["visualizations"], list):
                for viz in content["visualizations"]:
                    if isinstance(viz, str):
                        visualization_prompts.append(viz)
                    elif isinstance(viz, dict) and "description" in viz:
                        visualization_prompts.append(viz["description"])

            # Generate a default visualization if none specified
            if not visualization_prompts:
                # Create some generic visualizations based on the poster content
                if "methodology" in content:
                    methodology_text = " ".join(content["methodology"]) if isinstance(content["methodology"], list) else \
                    content["methodology"]
                    visualization_prompts.append(f"Diagram illustrating the research methodology: {methodology_text[:100]}")

                if "results" in content:
                    results_text = " ".join(content["results"]) if isinstance(content["results"], list) else content[
                        "results"]
                    visualization_prompts.append(f"Data visualization showing key results: {results_text[:100]}")

        # 使用API支持的尺寸
        # 根据错误信息，API支持的尺寸有: 1024x1024, 1152x896, 1216x832, 1344x768, 1536x640, 640x1536, 768x1344, 832x1216, 896x1152
        # 选择竖屏比例的尺寸
        api_supported_sizes = [
            (768, 1344),  # 竖屏比例
            (832, 1216),  # 竖屏比例
            (1024, 1024)  # 备选方案
        ]

        # 控制生成的图像数量 - 这里可以调整生成的海报数量
        max_images_to_generate = 1  # 设置为需要的数量，目前设置为1张

        # Generate each visualization
        for i, prompt in enumerate(visualization_prompts):
            # 限制生成的图像数量
            if i >= max_images_to_generate:
                break
                
            # 修改提示词，生成纯背景图像，不含文本
            enhanced_prompt = f"Academic poster background: {prompt}. Clean design, no text, pure background image only."

            success = False
            for width, height in api_supported_sizes:
                if success:
                    break

                try:
                    print(f"尝试生成尺寸为 {width}x{height} 的图像...")
                    response = requests.post(
                        self.api_url,
                        headers=self.headers,
                        json={
                            "text_prompts": [
                                {"text": enhanced_prompt, "weight": 1.0},
                                {"text": "text, words, labels, blurry, low quality, distorted", "weight": -1.0}
                            ],
                            "cfg_scale": 8,
                            "width": width,
                            "height": height,
                            "samples": 1,
                            "steps": 30,
                        }
                    )

                    if response.status_code == 200:
                        data = response.json()
                        for j, image in enumerate(data["artifacts"]):
                            # 保存为JPG格式
                            raw_image_path = os.path.join(config.TEMP_DIR, f"background_{i}_{j}_{uuid.uuid4()}.jpg")
                            with open(raw_image_path, "wb") as f:
                                f.write(base64.b64decode(image["base64"]))
                            
                            # 处理图像以确保竖版比例并填充页面
                            processed_image_path = self._process_image_for_display(raw_image_path, 1080, 1920)
                            image_paths.append(processed_image_path)
                            success = True
                        break
                    else:
                        print(f"生成尺寸 {width}x{height} 的图像失败: {response.status_code} - {response.text}")
                except Exception as e:
                    print(f"生成尺寸 {width}x{height} 的图像时发生异常: {str(e)}")

        # 如果没有成功生成任何图像，使用占位图
        if not image_paths:
            # 使用utils中的函数创建占位图
            from utils.image_utils import create_placeholder_image
            
            placeholder_path = os.path.join(config.TEMP_DIR, f"placeholder_{uuid.uuid4()}.png")
            placeholder_path = create_placeholder_image(
                placeholder_path, 
                width=1080, 
                height=1920, 
                text="图像生成失败，请重试"
            )
            image_paths.append(placeholder_path)

        return image_paths

    def _process_image_for_display(self, image_path: str, target_width: int = 1080, target_height: int = 1920) -> str:
        """处理图像以确保竖版比例并填充整个页面"""
        try:
            # 打开原始图像
            img = Image.open(image_path)
            
            # 创建竖版比例的新图像（白色背景）
            new_img = Image.new('RGB', (target_width, target_height), (255, 255, 255))
            
            # 计算调整后的尺寸，保持原始纵横比
            img_width, img_height = img.size
            
            # 使用填充模式而非适应模式，确保图像填满整个页面
            width_ratio = target_width / img_width
            height_ratio = target_height / img_height
            ratio = max(width_ratio, height_ratio)  # 使用最大比例确保填充满
            
            new_size = (int(img_width * ratio), int(img_height * ratio))
            
            # 调整原始图像大小
            resized_img = img.resize(new_size, Image.LANCZOS)
            
            # 计算裁剪区域，使图像居中
            left = (new_size[0] - target_width) // 2 if new_size[0] > target_width else 0
            top = (new_size[1] - target_height) // 2 if new_size[1] > target_height else 0
            right = left + target_width
            bottom = top + target_height
            
            # 裁剪图像
            if new_size[0] > target_width or new_size[1] > target_height:
                cropped_img = resized_img.crop((left, top, right, bottom))
                new_img = cropped_img
            else:
                # 如果调整后的图像小于目标尺寸，则居中放置
                position = ((target_width - new_size[0]) // 2, (target_height - new_size[1]) // 2)
                new_img.paste(resized_img, position)
            
            # 保存处理后的图像为JPG格式
            output_path = os.path.join(config.TEMP_DIR, f"processed_{os.path.basename(image_path)}")
            if not output_path.lower().endswith('.jpg'):
                output_path = output_path.rsplit('.', 1)[0] + '.jpg'
            new_img.save(output_path, 'JPEG', quality=95)
            
            return output_path
        except Exception as e:
            print(f"处理图像时出错: {str(e)}")
            return image_path  # 如果处理失败，返回原始图像路径